import torch

from transformers import BertTokenizer, BertModel

# 指定本地路径

model_path = "/root/ai_project/ai_doctor/model/bert-base-chinese"

# 加载tokenizer和模型
tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertModel.from_pretrained(model_path)

# # 使用tokenizer对文本进行编码
# text = "你好，世界！"
# inputs = tokenizer(text, return_tensors="pt")

# # 获取模型输出
# outputs = model(**inputs)
# print(outputs)


def get_bert_encode_for_single(text):
    """
    description: 使用bert-chinese编码中文文本
    :param text: 要进行编码的文本
    :return: 使用bert编码后的文本张量表示
    """
    # 首先使用字符映射器对每个汉字进行映射
    # 这里需要注意，bert的tokenizer映射后会为结果前后添加开始和结束标记即101和102 
    # 这对于多段文本的编码是有意义的，但在我们这里没有意义，因此使用[1:-1]对头和尾进行切片
    indexed_tokens = tokenizer.encode(text)[1:-1]
    # 之后将列表结构转化为tensor
    tokens_tensor = torch.LongTensor([indexed_tokens])
    #print('tokens_tensor:', tokens_tensor)
    # 使模型不自动计算梯度
    with torch.no_grad():
        # 调用模型获得隐层输出
        encoded_layers = model(tokens_tensor)
    # 输出的隐层是一个三维张量，最外层一维是1, 我们使用[0]降去它。
    #print('encoded_layers:', encoded_layers)
    encoded_layers = encoded_layers[0]
    return encoded_layers

# text = "你好，周杰伦"
# text = "你好，世界！"

# outputs = get_bert_encode_for_single(text)
# print(outputs)
# print(outputs.shape)
